package top.bug.blog.aop;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.util.DigestUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import top.bug.blog.annotation.RateLimit;
import top.bug.blog.common.CustomRateLimiter;
import top.bug.blog.enums.ErrorCode;
import top.bug.blog.enums.RateLimitEnum;
import top.bug.blog.exception.BussinessException;
import top.bug.blog.utils.IpUtils;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * @author <a href="https://github.com/jiusi66">九思.</a>
 * @email 10837907@qq.com
 * @from <a href="https://www.52bug.top">递归茶馆-www.52bug.top</a>
 * @description 接口限流拦截器-支持阶梯惩罚机制
 */
@Aspect
@Component
@Slf4j
public class RateLimitAspect {
    private static final String RATE_LIMIT_KEY_PREFIX = "rate_limit:";

    /**
     * 最大屏蔽时长限制,避免无限增长 - 默认 2小时
     */
    private static final int MAX_BLOCK_DURATION = 7200;

    /**
     * 惩罚因子，默认0.5
     */
    private static final double PENALTY_FACTOR = 0.5;


    /**
     * 使用Caffeine缓存存储计数器（自动过期）
     */
    private final Cache<String, CustomRateLimiter> counterCache = Caffeine.newBuilder()
            .expireAfterAccess(2, TimeUnit.HOURS)
            .maximumSize(10_000)
            .build();

    /**
     * 使用单独缓存存储违规记录
     */
    private final Cache<String, Integer> violationCache = Caffeine.newBuilder()
            // 违规记录保留24小时
            .expireAfterWrite(24, TimeUnit.HOURS)
            .maximumSize(5_000)
            .build();

    @Around("@annotation(rateLimit)")
    public Object checkRateLimit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder
                .currentRequestAttributes()).getRequest();

        // 获取实际的限流参数（优先使用枚举配置）
        String effectiveKey;
        int effectiveMaxAttempts;
        int effectiveTimeWindow;
        int effectiveBlockDuration;

        if (rateLimit.strategyEnum() != RateLimitEnum.UNSPECIFIED_LIMIT) {
            // 使用枚举配置
            effectiveKey = rateLimit.strategyEnum().getKey();
            effectiveMaxAttempts = rateLimit.strategyEnum().getMaxAttempts();
            effectiveTimeWindow = rateLimit.strategyEnum().getTimeWindow();
            effectiveBlockDuration = rateLimit.strategyEnum().getBlockDuration();
        } else {
            // 使用注解直接配置的参数
            effectiveKey = rateLimit.key();
            effectiveMaxAttempts = rateLimit.maxAttempts();
            effectiveTimeWindow = rateLimit.timeWindow();
            effectiveBlockDuration = rateLimit.blockDuration();
        }

        // 1. 获取客户端真实IP（处理反向代理情况）
        String ip = IpUtils.getClientIp(request);

        // 2. 构建缓存键（方法签名+IP）
        String cacheKey = buildCacheKey(joinPoint, effectiveKey, ip);

        // 3. 获取当前计数器
        CustomRateLimiter limiter = counterCache.get(cacheKey, k ->
                new CustomRateLimiter(0, System.currentTimeMillis() + effectiveTimeWindow * 1000L, 0)
        );

        // 4. 检查是否在屏蔽期
        if (System.currentTimeMillis() < limiter.getBlockExpireTime()) {
            throw new BussinessException(ErrorCode.OPERATE_TOO_FREQUENTLY);
        }

        // 5. 检查时间窗口重置
        if (System.currentTimeMillis() > limiter.getResetTime()) {
            limiter.setCounter(0);
            limiter.setResetTime(System.currentTimeMillis() + effectiveTimeWindow * 1000L);
        }

        // 6. 增加计数
        limiter.setCounter(limiter.getCounter() + 1);
        counterCache.put(cacheKey, limiter);

        // 7. 超限处理
        if (limiter.getCounter() > effectiveMaxAttempts) {
            // 获取该IP的历史违规次数
            Integer violationCount = violationCache.get(ip, k -> 0);
            // 阶梯惩罚计算
            double penaltyFactor = calculatePenaltyFactor(violationCount);
            int dynamicBlockTime = applyPenalty(effectiveBlockDuration, penaltyFactor);

            // 更新违规记录
            violationCache.put(ip, violationCount + 1);

            // 设置动态屏蔽时间
            limiter.setBlockExpireTime(System.currentTimeMillis() + dynamicBlockTime * 1000L);
            counterCache.put(cacheKey, limiter);
            // 记录日志
            log.warn("IP[{}]触发阶梯限流 - 违规次数: {} - 屏蔽时长: {}秒", ip, violationCount + 1, dynamicBlockTime);

            throw new BussinessException(ErrorCode.OPERATE_TOO_FREQUENTLY,
                    "操作过于频繁，请" + dynamicBlockTime + "秒后重试");
        }

        // 8. 放行请求
        return joinPoint.proceed();
    }

    /**
     * 构建缓存键
     * 格式：rate_limit:[自定义key或方法名]:[IP]
     */
    private String buildCacheKey(ProceedingJoinPoint joinPoint, String customKey, String ip) {
        if (StringUtils.hasText(customKey)) {
            return RATE_LIMIT_KEY_PREFIX + customKey + ":" + ip;
        }
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        String methodName = method.getName();
        // 获取方法名全名
        String methodFullName = method.getDeclaringClass().getName() + "." + methodName;
        // 取MD5
        String sign = DigestUtils.md5DigestAsHex(methodFullName.getBytes());
        return RATE_LIMIT_KEY_PREFIX + sign + ":" + ip;
    }

    /**
     * 计算惩罚因子
     */
    private double calculatePenaltyFactor(int violationCount) {
        // 每违规一次，惩罚因子增加0.5
        return 1.0 + PENALTY_FACTOR * violationCount;
    }

    /**
     * 应用惩罚规则
     */
    private int applyPenalty(int baseDuration, double factor) {
        int dynamicDuration = (int) (baseDuration * factor);
        return Math.min(dynamicDuration, MAX_BLOCK_DURATION);
    }
}